#!/usr/bin/env python
import rospy
from geometry_msgs.msg import PoseStamped
from mavros_msgs.msg import State, CommandBool, SetMode, PositionTarget
from sensor_msgs.msg import NavSatFix
from std_msgs.msg import Float64
import math

current_state = State()
local_pos = PoseStamped()
last_local_pos = PoseStamped()
global_pos = NavSatFix()

R_EARTH = 6371000  # Earth radius in meters

def state_cb(msg):
    global current_state
    current_state = msg

def local_pos_callback(msg):
    global last_local_pos, local_pos
    last_local_pos = local_pos
    local_pos = msg

def global_pos_callback(msg):
    global global_pos
    global_pos = msg

def global_to_local_end(lat1, lon1, lat2, lon2):
    dLat = (lat2 - lat1) * math.pi / 180.0
    dLon = (lon2 - lon1) * math.pi / 180.0
    latAvg = (lat1 + lat2) / 2.0 * math.pi / 180.0
    cosLat = math.cos(latAvg)
    dx = R_EARTH * dLon * cosLat
    dy = R_EARTH * dLat
    return dx, dy

def main():
    rospy.init_node('test_control_node', anonymous=True)
    rate = rospy.Rate(20.0)

    state_sub = rospy.Subscriber('/mavros/state', State, state_cb)
    local_pos_sub = rospy.Subscriber('/mavros/local_position/pose', PoseStamped, local_pos_callback)
    global_pos_sub = rospy.Subscriber('/mavros/global_position/global', NavSatFix, global_pos_callback)
    local_pos_pub = rospy.Publisher('/mavros/setpoint_raw/local', PositionTarget, queue_size=10)
    arming_client = rospy.ServiceProxy('/mavros/cmd/arming', CommandBool)
    set_mode_client = rospy.ServiceProxy('/mavros/set_mode', SetMode)

    while not rospy.is_shutdown() and not current_state.connected:
        rate.sleep()

    my_px4_mode = "OFFBOARD"
    offb_set_mode = SetMode()
    offb_set_mode.request.custom_mode = my_px4_mode

    arm_cmd = CommandBool()
    arm_cmd.request.value = True
    last_request_time = rospy.Time.now()
    last_local_pos_time = rospy.Time.now()

    pos_errorx = [Float64() for _ in range(3)]
    pos_errory = [Float64() for _ in range(3)]
    pos_errorz = [Float64() for _ in range(3)]

    for i in range(3):
        pos_errorx[i].data = 10000
        pos_errory[i].data = 10000
        pos_errorz[i].data = 10000

    target_pose = PositionTarget()
    target_pose.header.frame_id = "base_footprint"
    target_pose.coordinate_frame = PositionTarget.FRAME_LOCAL_NED
    target_pose.type_mask = PositionTarget.IGNORE_AFX | PositionTarget.IGNORE_AFY | PositionTarget.IGNORE_AFZ | PositionTarget.IGNORE_YAW_RATE
    target_pose.position.x = 0.0
    target_pose.position.y = 0.0
    target_pose.position.z = 3.0
    target_pose.yaw = 0.0

    target_lat = 47.397743
    target_lon = 8.545585
    target_alt = 580.443928

    dx, dy = global_to_local_end(global_pos.latitude, global_pos.longitude, target_lat, target_lon)
    dz = target_alt - global_pos.altitude
    dz = max(min(dz, 30), -50)
    dx = max(min(dx, 100), -100)
    dy = max(min(dy, 100), -100)

    rospy.loginfo(f"lat: {global_pos.latitude}, lon: {global_pos.longitude}, alt: {global_pos.altitude}")
    rospy.loginfo(f"target_lat: {target_lat}, target_lon: {target_lon}, target_alt: {target_alt}")
    rospy.loginfo(f"dx: {dx}, dy: {dy}, dz: {dz}")

    for i in range(100):
        if rospy.is_shutdown():
            break
        local_pos_pub.publish(target_pose)
        rate.sleep()

    target_pose.position.x = local_pos.pose.position.x + dx
    target_pose.position.y = local_pos.pose.position.y + dy
    target_pose.position.z = local_pos.pose.position.z + dz

    if current_state.mode != my_px4_mode:
        if set_mode_client(offb_set_mode).mode_sent:
            rospy.loginfo(f"{my_px4_mode} enabled")
        last_request_time = rospy.Time.now()

    if not current_state.armed:
        if arming_client(arm_cmd).success:
            rospy.loginfo("Vehicle armed")
        last_request_time = rospy.Time.now()

    while not rospy.is_shutdown():
        target_pose.header.stamp = rospy.Time.now()
        local_pos_pub.publish(target_pose)

        if rospy.Time.now() - last_local_pos_time > rospy.Duration(0.8):
            pos_errorx[2].data = pos_errorx[1].data
            pos_errory[2].data = pos_errory[1].data
            pos_errorz[2].data = pos_errorz[1].data

            pos_errorx[1].data = pos_errorx[0].data
            pos_errory[1].data = pos_errory[0].data
            pos_errorz[1].data = pos_errorz[0].data

            pos_errorx[0].data = local_pos.pose.position.x - target_pose.position.x
            pos_errory[0].data = local_pos.pose.position.y - target_pose.position.y
            pos_errorz[0].data = local_pos.pose.position.z - target_pose.position.z

            rospy.logwarn(f"local_pos.x, y, z: {local_pos.pose.position.x}, {local_pos.pose.position.y}, {local_pos.pose.position.z}")
            rospy.loginfo(f"pos_errorx[0].data: {pos_errorx[0].data}, pos_errory[0].data: {pos_errory[0].data}, pos_errorz[0].data: {pos_errorz[0].data}")

            all_in_range = all(-0.6 < err.data < 0.6 for err in [pos_errorx[0], pos_errory[0], pos_errorz[0]])

            if all_in_range:
                rospy.loginfo("Position is within range")
                hold_mode = SetMode()
                hold_mode.request.custom_mode = "AUTO.LOITER"
                if set_mode_client(hold_mode).mode_sent:
                    rospy.loginfo("Switched to LOITER mode. The drone will maintain its position.")
                break

            last_local_pos_time = rospy.Time.now()

        rate.sleep()

if __name__ == '__main__':
    try:
        main()
    except rospy.ROSInterruptException:
        pass
